#
# Copyright 2013 Antoine Drouin (poinix@gmail.com)
#
# This file is part of PAT.
#
#    PAT is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    PAT is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with PAT.  If not, see <http://www.gnu.org/licenses/>.
#

"""
 This is a 6dof synthesis model for an RC helicopter as describded
 in the paper "Autonomous Autorotation of an RC Helicopter" by
 Abbeel, Coates, Hunter and Ng
"""

import math
import numpy as np
import scipy
from scipy import optimize
import matplotlib.pyplot as plt

import pat.dynamic_model as dm
import pat.algebra as pa
import pat.utils   as pu

"""
State Vector Components
"""
sv_x     =  0     # position in an euclidian space
sv_y     =  1     #
sv_z     =  2     # 
sv_u     =  3     # velocity in body frame
sv_v     =  4     #
sv_w     =  5     #
sv_phi   =  6     # attitude
sv_theta =  7     #
sv_psi   =  8     # 
sv_p     =  9     # rotational velocities
sv_q     = 10     #
sv_r     = 11     #
sv_om    = 12     # rotor speed
sv_size  = 13

"""
Input Vector Components
"""
iv_lat   = 0      # lateral cyclic
iv_lon   = 1      # longitudinal cyclic
iv_rud   = 2      # tail rotor
iv_col   = 3      # collective pitch
iv_thr   = 4      # throttle
iv_size  = 5


def dyn(X, t, U, P):
    """
    Dynamic model
    """
    X_vel_body  = X[sv_u:sv_w+1]
    X_euler     = X[sv_phi:sv_psi+1]
    X_rvel_body = X[sv_p:sv_r+1]
    vel_uv = math.sqrt(X[sv_u]**2+X[sv_v]**2)
    earth_to_body = pa.rmat_of_euler(X_euler)

    Xd = np.zeros(sv_size)

    Xd[sv_x:sv_z+1] = np.dot(np.transpose(earth_to_body), X_vel_body)

    g_body = np.dot(earth_to_body, [0, 0, 9.81])
    a = [P.Cu*X[sv_u],
         P.Cv*X[sv_v],
         np.dot(P.Cw,[1., X[sv_w], U[iv_col]*X[sv_om], vel_uv])]
    Xd[sv_u:sv_w+1] = np.cross(X_rvel_body, X_vel_body) + g_body + a

    Xd[sv_phi:sv_psi+1] = pa.euler_derivatives(X_euler, X_rvel_body)

    Xd[sv_p:sv_r+1] = [np.dot(P.Cp, [1., X[sv_p], U[iv_lat]*X[sv_om]]),
                       np.dot(P.Cq, [1., X[sv_q], U[iv_lon]*X[sv_om]]),
                       np.dot(P.Cr, [1., X[sv_r], U[iv_rud]*X[sv_om]])]

    a = [1., X[sv_om], U[iv_col], X[sv_w], vel_uv, U[iv_lat]**2+U[iv_lon]**2, U[iv_thr]]
    Xd[sv_om] = np.dot(P.Co,a)
    return Xd


def trim(P, args, debug=False):
    if args <>None:
        om = args['om']
        vel = args['vel']
        gamma = 0.
    else:
        om = pu.rps_of_rpm(1700.)
        vel = 0.
        gamma = 0.

    vel_earth = vel*np.array([math.cos(gamma), 0., -math.sin(gamma)])
    print "vel_earth", vel_earth
    theta0 = math.atan2(P.Cu*vel_earth[0], 9.81+P.Cu*vel_earth[2])
    print "theta0: ", pu.deg_of_rad(theta0), "deg"
    def err_fun(p):
        earth_to_body = pa.rmat_of_euler([p[0]/1000., theta0, 0.])
        u, v, w = np.dot(earth_to_body, vel_earth)
#        print "theta {:f} u{:f} w{:f}".format(pu.deg_of_rad(p[1]), u, w)
        X = [0., 0., 0.,  u, v, w,  p[0]/1000., p[1]/1000., 0.,  0., 0., 0., om]
        Xdot = dyn(X, 0., p[1:], P)
        Xdot_ref = [vel_earth[0], vel_earth[1], vel_earth[2],   0., 0., 0.,   0., 0., 0.,  0., 0., 0.,  0.]
        return np.linalg.norm(Xdot - Xdot_ref)

    #    phi     lat    lon   rud  col, thr 
    p = [ 0.,    0.1,  -0.15,  0.,  1.,  1.]
    p = scipy.optimize.fmin_powell(err_fun, p, disp=debug, ftol=1e-9)
    print "phi:", pu.deg_of_rad(p[0]/1000.)
    earth_to_body = pa.rmat_of_euler([p[0]/1000., p[1]/1000., 0.])
    u, v, w = np.dot(earth_to_body, vel_earth) 
    Xe = [0., 0., -2.,  u, v, w,  p[0]/1000., theta0, 0.,  0., 0., 0.,  om]
    Ue = p[1:]
    return Xe, Ue
    
    


class Param():
    def __init__(self, filename=None):
        self.Cu = np.array([-0.05])
        self.Cv = np.array([-0.06])
        self.Cw = np.array([ -0.47, -1.42,  -0.01, -0.15])
        self.Cp = np.array([ -1.46, -5.74,   0.02])
        self.Cq = np.array([ -0.23, -5.32,  -0.01])
        self.Cr = np.array([  0.52, -5.43,   0.02])
        self.Co = np.array([106.85, -0.23, -68.53, 22.79, 2.11, -6.10, 50.])
        self.name = "default"

    def __str__(self):
        str = """Cu:{:s}\nCv:{:s}\nCw:{:s}\nCp:{:s}\nCq:{:s}\nCr:{:s}\nCo:{:s}
        """.format(self.Cu, self.Cv, self.Cw, self.Cp, self.Cq, self.Cr, self.Co)
        return str




class DynamicModel(dm.DynamicModel):

    sv_x = sv_x
    sv_size = sv_size
    iv_size = iv_size
    
    def __init__(self, param=None):
        dm.DynamicModel.__init__(self)
        self.P = Param()
        self.X = np.zeros(sv_size)

    def name(self):
        return "Rotorcraft Python Synthesis ({:s})".format(self.P.name)

    def reset(self, X0=None):
        if X0<>None: self.X = X0
        else: self.X = np.array([0., 0., -2.,  0., 0., 0.,  0., 0., 0.,  0., 0., 0., pu.rps_of_rpm(1700)])
        return self.X

    def trim(self, args=None, debug=False):
        return trim(self.P, args, debug)

    def get_jacobian(self, Xe, Ue): return pu.num_jacobian(Xe, Ue, self.P, dyn)

    def run(self, dt, U):
        foo, self.X = scipy.integrate.odeint(dyn, self.X, [0, dt], args=(U, self.P, ))
        return self.X

    def param(self):
        return str(self.P)

    def state_SixDOFfEuclidianEuler(self):
        return self.X

    def state_str(self):
        return """
        pos: {:-.2f}, {:-.2f}, {:-.2f} m
        vel: {:-.2f}, {:-.2f}, {:-.2f} m/s, 
        att: {:-.2f}, {:-.2f}, {:-.2f} deg
        om : {:-.2f} rpm
        """.format(self.X[sv_x], self.X[sv_y], self.X[sv_z],
        self.X[sv_v], pu.deg_of_rad(self.X[sv_u]), pu.deg_of_rad(self.X[sv_v]),
        pu.deg_of_rad(self.X[sv_phi]), pu.deg_of_rad(self.X[sv_theta]), pu.deg_of_rad(self.X[sv_psi]),
        pu.rpm_of_rps(self.X[sv_om]))  


    def plot_trajectory(self, time, X, U=None, Xref=None, figure=None,
                        window_title="Trajectory", legend=None, filename=None):

        margins=(0.04, 0.05, 0.98, 0.96, 0.20, 0.34)
        figure = pu.prepare_fig(figure, window_title, figsize=(20.48, 10.24), margins=margins)

        plots = [("x", "m", X[:,sv_x]), ("y", "m", X[:,sv_y]), ("z", "m", X[:,sv_z]),
                 ("u", "m/s", X[:,sv_u]), ("v", "m/s", X[:,sv_v]), ("w", "m/s", X[:,sv_w]),
                 ("$\phi$",    "deg", pu.deg_of_rad(X[:,sv_phi])),
                 ("$\\theta$", "deg", pu.deg_of_rad(X[:,sv_theta])),
                 ("$\\psi$",   "deg", pu.deg_of_rad(X[:,sv_psi])),
                 ("$p$",     "deg/s", pu.deg_of_rad(X[:,sv_p])),
                 ("$q$",     "deg/s", pu.deg_of_rad(X[:,sv_q])),
                 ("$r$",     "deg/s", pu.deg_of_rad(X[:,sv_r])),
                 ("$\omega$","rpm",   pu.rpm_of_rps(X[:,sv_om]))]

        if Xref<>None:
            refs = [Xref[:,sv_x], Xref[:,sv_y], Xref[:,sv_z], Xref[:,sv_u], Xref[:,sv_v], Xref[:,sv_w],
                    pu.deg_of_rad(Xref[:,sv_phi]), pu.deg_of_rad(Xref[:,sv_theta]),
                    pu.deg_of_rad(Xref[:,sv_psi]), pu.deg_of_rad(Xref[:,sv_p]),
                    pu.deg_of_rad(Xref[:,sv_q]), pu.deg_of_rad(Xref[:,sv_r]), pu.rpm_of_rps(Xref[:,sv_om])]
    
        for i, (title, ylab, data) in enumerate(plots):
            ax = plt.subplot(5, 3, i+1)
            plt.plot(time, data)
            if Xref<>None: plt.plot(time, refs[i])
            pu.decorate(ax, title=title, ylab=ylab)
            
        if legend<>None:
            plt.legend(legend, loc='best')

        if U<>None:
            ax = plt.subplot(5, 3, 14)
            plt.plot(time, U[:, iv_lat:iv_rud+1])
            plt.legend(["lat", "lon", "rud"], loc='best')
            ax = plt.subplot(5, 3, 15)
            plt.plot(time, U[:, iv_col:iv_thr+1])
            plt.legend(["col", "throttle"], loc='best')

        pu.save_if(filename)

        return figure
        
